package sketch.ounit.fuzz;

import java.io.File;
import java.lang.reflect.Method;
import java.util.Collection;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import org.eclipse.jdt.core.dom.Expression;
import org.eclipse.jdt.core.dom.ImportDeclaration;
import org.eclipse.jdt.core.dom.MethodDeclaration;
import org.eclipse.jdt.core.dom.MethodInvocation;
import org.eclipse.jdt.core.dom.Modifier;

import randoop.main.GenInputsAbstract;
import randoop.util.Timer;
import sketch.ast.ASTUtils;
import sketch.generator.JUnitTestCreator;
import sketch.instrumenter.Tracer;
import sketch.main.CodeToolkit;
import sketch.main.SketchTestOptions;
import sketch.specs.annotation.TestSketch;
import sketch.util.Checker;
import sketch.util.Globals;
import sketch.util.IntermediateFiles;
import sketch.util.SketchUtil;

/**
 * This class is a new implementation for sketch-guided random test generation.
 * Please use this instead of the old implementation of SketchGuidedRandomTestGenerator.
 * 
 * XXX TODO possible limitation, the JUnit only supports the simplest form
 * no setUp / tearDown, all methods are unit test methods
 * 
 * Further change should be made to the checkJUnitMethod() method
 * */
public class ImprovedSketchGuidedRandomTestGeneration {
	
	/**
	 * All following fields are initialized in constructors
	 * */
	public final String sketchPath;
	public final String junit_class;
	public final int random_time_limit;
	public final int fuzz_time_limit;
	private final Map<MethodDeclaration, List<Expression>> receiver_map;
	private final Map<MethodDeclaration, List<MethodInvocation>> invoke_map;
	private final Map<MethodDeclaration, List<Expression[]>> arg_expr_map;
	private final RandomPlacesExtractor extractor;
	private final Collection<Class<?>> classes;
	
	/**
	 * Several configuration options
	 * */
	private String tmp = "./tmp"; //for storing compiled junit_class, must be in classpath
	//private String output_dir = "./output"; //for storing the final output
	private String junit_method_append = "_guided";
	private String junit_class_prefix = "RandoopTest";
	private int maxTestNum = 20000;
	private int testPerFile = randoop.main.GenInputsAbstract.testsperfile; //default setting here
	private String outputPackageName = null;
	private boolean outputRandoopTest = false;
	
	/**
	 * Default constructor
	 * */
	public ImprovedSketchGuidedRandomTestGeneration(String sketchPath, String junitClass, int randomtime, int fuzztime) {
		Checker.checkNull(sketchPath, "The sketchPath could not be null.");
		Checker.checkNull(junitClass, "The junitClass could not be null.");
		Checker.checkTrue(randomtime >= 0, "random time must be >= 0, the input: " + randomtime + " is not valid.");
		Checker.checkTrue(fuzztime >= 0, "fuzztime must be >= 0, the input: " + fuzztime + " is not valid.");
		this.sketchPath = sketchPath;
		this.junit_class = junitClass;
		this.random_time_limit = randomtime;
		this.fuzz_time_limit = fuzztime;
		//initialize the map
		this.extractor = new RandomPlacesExtractor(this.sketchPath);
		this.receiver_map = extractor.getExprMap();
		this.invoke_map = extractor.getInvokeMap();
		this.arg_expr_map = extractor.getArgExprMap();
		//find and load all classes from classpath
		classes = this.init_classes();
	}
	
	//all automated generated settings
	public void setTmp(String tmp) {
		this.tmp = tmp;
	}
//	public void setOutput_dir(String output_dir) {
//		this.output_dir = output_dir;
//	}
	public void setJunit_method_append(String junit_method_append) {
		this.junit_method_append = junit_method_append;
	}
	public void setJunit_class_prefix(String junit_class_prefix) {
		this.junit_class_prefix = junit_class_prefix;
	}
	public void setMaxTestNum(int maxTestNum) {
		Checker.checkTrue(maxTestNum > 0, "The maxTestNum: " + maxTestNum + " is not valid");
		this.maxTestNum = maxTestNum;
	}
	public void setTestPerFile(int testperfile) {
		Checker.checkTrue(testperfile > 0, "The testperfile: " + testperfile + " is not valid");
		this.testPerFile = testperfile;
	}
	public void setOutputPackageName(String outputPackageName) {
		this.outputPackageName = outputPackageName;
	}
	public void setOutputRandoopTest(boolean outputRandoopTest) {
		this.outputRandoopTest = outputRandoopTest;
	}
	
	/**************************************************************************
	 * The main entry for test generation
	 * The actual ONLY method for invocation.
	 * *************************************************************************/
	public List<File> generate_tests() {
		if(receiver_map.isEmpty()) {
			if(SketchTestOptions.verbose) {
			    System.out.println("No random values are specified, can not generate tests. " +
			    		"Just return untouched sketches.");
			}
			Checker.checkTrue(invoke_map.isEmpty(), "The invoke map must be empty.");
			//it should return the original sketch file without any changes
			List<File> untouchFiles = new LinkedList<File>();
			untouchFiles.add(new File(this.sketchPath));
			return untouchFiles;
		}
		Checker.checkTrue(receiver_map.size() == invoke_map.size(), "Their size are not the same!");
		Checker.checkTrue(receiver_map.size() == arg_expr_map.size(), "Their size are not the same!");
		//compile the given test sketch
		this.compileSketch();
		//generate tests, and output them to disk
		return this.sketch_guided_random_gen();
	}
	
	/*****************************************************************************
	 * All the following methods are all private, only accessible from this class
	 * ***************************************************************************/
	
	/**
	 * Parses the sketch to get all class declarations for random test generation.
	 * */
	private Collection<Class<?>> init_classes() {
		Collection<Class<?>> clz = new HashSet<Class<?>>();
		for(List<MethodInvocation> mlist : this.invoke_map.values()) {
			for(MethodInvocation invoke : mlist) {
				List<Expression> args = invoke.arguments();
				for(Expression arg : args) {
					String argStr = arg.toString();
					if(argStr.endsWith(".class")) {
						String className = argStr.substring(0, argStr.length() - ".class".length());
						try {
							Class<?> c = Class.forName(className);
							clz.add(c);
						} catch (Exception e) {
							System.err.println("class: " + className + " can not be found.");
						}
					}
				}
			}
		}
		return clz;
	}
	
	/**
	 * Compiles the sketch, and output the class file to the tmp dir
	 * */
	private void compileSketch() {
		String packageName = extractor.getPackageName();
		String compile_dir = tmp + Globals.fileSep + SketchUtil.transClassNameDotToSlash(packageName);
		SketchUtil.createDir(compile_dir);
		//compile the code
		List<File> files = new LinkedList<File>();
		File sketchFile = new File(this.sketchPath);
		files.add(sketchFile);
		//CodeToolkit.compile(files);
		CodeToolkit.compileToDir(files, tmp);
		//add to intermediate files
		IntermediateFiles.addAllFilesAsIntermediates(compile_dir, ".class");
	}
	
	/**
	 * Generates tests randomly guided by a test sketch
	 * */
	private List<File> sketch_guided_random_gen() {
		//first step: create object pool
		SharedSequencePool.init_pool(this.random_time_limit, (Class<?>[])this.classes.toArray(new Class<?>[0]));
		//then guided test generation
		Map<MethodDeclaration, List<List<AssembledSequence<?>>>>  generated_method_sequences = execute_junit_to_create_tests();
		//generate all tests
		return this.create_and_output_tests(generated_method_sequences);
	}
	
	/**
	 * The core algorithm of sketch-guided test generation. It repeatedly executes a given junit sketch.
	 * Each execution produces a new created test.
	 * */
	private Map<MethodDeclaration, List<List<AssembledSequence<?>>>> execute_junit_to_create_tests() {
		//the return map
		Map<MethodDeclaration, List<List<AssembledSequence<?>>>> seq_map =
			new LinkedHashMap<MethodDeclaration, List<List<AssembledSequence<?>>>>();
		//first, we load the test sketch class
		Class<?> clazz = CodeToolkit.loadClassFromDir(tmp, this.junit_class);
		this.checkJUnitMethod(clazz);
		
		//create an instance of a given class
		Object junit_test_obj = null;
		try {
			junit_test_obj = clazz.newInstance();
		} catch(Exception e) {
			throw new RuntimeException(e);
		}
		
		Method setUp = this.getPublicVoidNoArgMethod("setUp", clazz);
		Method tearDown = this.getPublicVoidNoArgMethod("tearDown", clazz);
		
		//establish the method declaration to real method mapping
	    Map<MethodDeclaration, List<List<AssembledSequence<?>>>> unique_sequences
	       = new LinkedHashMap<MethodDeclaration, List<List<AssembledSequence<?>>>>();
		Map<MethodDeclaration, Method> method_map = new LinkedHashMap<MethodDeclaration, Method>();
		for(MethodDeclaration method : this.invoke_map.keySet()) {
			//create an empty list
			List<List<AssembledSequence<?>>> unique_pools = new LinkedList<List<AssembledSequence<?>>>();
			unique_sequences.put(method, unique_pools);
			Method exec_junit_method = null;
			try {
				exec_junit_method = clazz.getMethod(method.getName().toString());
			} catch (Exception e) {
				throw new RuntimeException(e);
			}
			method_map.put(method, exec_junit_method);
		}
		
		//start to generate tests
		Timer t = new Timer();
		int count = 0;
		t.startTiming();
		while(count < this.maxTestNum && t.getTimeElapsedMillis() < 1000*this.fuzz_time_limit) {
			//execute each method one by one
			for(MethodDeclaration method : this.invoke_map.keySet()) {
				//reset the pool, and start a new round
			    SharedSequencePool.startNewRound();
				//get the method, and execute that reflectively
				Method exec_junit_method = method_map.get(method);
				exec_junit_method.setAccessible(true);
				Checker.checkTrue(exec_junit_method.isAccessible() && exec_junit_method.getReturnType().equals(void.class)
						&& !Modifier.isStatic(exec_junit_method.getModifiers()), "method: " + exec_junit_method + " is not a junit method.");
				
				try {
					//this is somehow tricky, need to prevent the Tracer record
					//before doing random test generation
					Tracer.turnOff();
					//System.out.println("start invoke...");
					if(setUp != null) {
						setUp.invoke(junit_test_obj);
					}
				    exec_junit_method.invoke(junit_test_obj);
				    if(tearDown != null) {
				    	tearDown.invoke(junit_test_obj);
				    }
				    Tracer.turnOn();
				} catch (Exception e) {
					//skip the exception-throwing cases
					Tracer.turnOn();
					continue;
				}
				//get the generated sequence
				List<AssembledSequence<?>> seq_list = SharedSequencePool.getGeneratedSequence();
				unique_sequences.get(method).add(seq_list);
				count++;
			}
		}
		seq_map.putAll(unique_sequences);
		t.stopTiming();
		
		return seq_map;
	}
	
	/**
	 * Creates a set of JUnit tests based on the produced sequences.
	 * #XXX note, it only return files created by sketch-guided random testing, randoop created
	 * files are not returned at all.
	 * */
	private List<File> create_and_output_tests(Map<MethodDeclaration, List<List<AssembledSequence<?>>>>  generated_method_sequences) {
		if(SketchTestOptions.verbose) {
		    System.out.println("Sketched methods: " + generated_method_sequences.keySet().size());
		}
		String othercode = extractor.getOtherMethodCode();
		List<File> files = new LinkedList<File>();
		//generate compilation unit for each test sketch method
		for(MethodDeclaration method : generated_method_sequences.keySet()) {
			List<String> outputTestMethods = new LinkedList<String>();
			//the string builder
			//StringBuilder sb = new StringBuilder();
			//all expressions, invocations, and argument expressions for the given test sketch
			List<Expression> expressions = this.receiver_map.get(method);
			List<MethodInvocation> invokes = this.invoke_map.get(method);
			List<Expression[]> arg_expressions = this.arg_expr_map.get(method);
			Checker.checkTrue(expressions.size() == invokes.size(), "The expression size: "
					+ expressions.size() + ", the invoke size: " + invokes.size());
			//get the generated sequences
			List<List<AssembledSequence<?>>> lists = generated_method_sequences.get(method);
			this.checkEqual(lists);
			//create JUnit tests
			int generated_num = 0;
			for(List<AssembledSequence<?>> list : lists) {
			   List<String> sequences_str = new LinkedList<String>();
			   List<String> result_vars = new LinkedList<String>();
			   //do the renaming here
			   int length_of_created_sequence = 0;
			   int arg_name_count = 0;
			   for(int i = 0; i < list.size(); i++) {
				   Expression[] exprs = arg_expressions.get(i);
				   String[] arg_names = new String[exprs.length];
				   String[] arg_exprs = new String[exprs.length];
				   for(int exp_index = 0; exp_index < exprs.length; exp_index++) {
					   arg_names[exp_index] = "arg_var" + (arg_name_count++);
					   arg_exprs[exp_index] = exprs[exp_index].toString();
				   }
				   
				   AssembledSequence<?> a_sequence = list.get(i);
				   String code = a_sequence.toCodeString(length_of_created_sequence, expressions.get(i).toString(), arg_names, arg_exprs);
				   result_vars.add(a_sequence.getResultVar(length_of_created_sequence));
				   sequences_str.add(code);
				   //update the length
				   length_of_created_sequence += list.get(i).size();
			   }
			   MethodDeclaration copy = ASTUtils.deepClone(method);
			   //create a new JUnit method, setting the name
			   copy.getName().setIdentifier(method.getName().getIdentifier() + this.junit_method_append + (generated_num++));
			   ASTUtils.removeMethodAnnotation(copy, TestSketch.class);
			   RandomFuzzInstantiator instantiator = new RandomFuzzInstantiator(copy, invokes, sequences_str, result_vars);
			   //create tests only when getting valid sequences
			   if(invokes.size() == sequences_str.size()) {
			     String method_str = instantiator.getResult();
//			     sb.append(method_str);
//			     sb.append(Globals.lineSep);
			     outputTestMethods.add(method_str + Globals.lineSep);
			   } else {
				   if(SketchTestOptions.verbose) {
				       System.out.println("Discard this sequence, invokes.size() != sequences_str.size(), invoke size: "
						   + invokes.size() + ", sequence string size: " + sequences_str.size());
				   }
			   }
			}
			if(SketchTestOptions.verbose) {
			    System.out.println("Generated num in lists: " + lists.size() + " for method: " + method.getName());
			}
			
			//write JUnit test
			int testNum = (outputTestMethods.size() / this.testPerFile) + (outputTestMethods.size() % this.testPerFile == 0 ? 0 : 1);
			String output_name = this.junit_class_prefix + method.getName().getIdentifier() + "GenTest";
			int startIndex = 0;
			for(int i = 0; i < testNum; i++) {
				String fileName = output_name + i;
				StringBuilder sb = new StringBuilder();
				//append all Java
				sb.append(othercode);
				
				for(int ci = 0;  ci < this.testPerFile; ci++) {
					int test_index = startIndex + ci;
					if(test_index >= outputTestMethods.size()) {
						break;
					}
					sb.append(outputTestMethods.get(test_index));
				}
				
			    File f = JUnitTestCreator.createJUnitTest(sb.toString(), SketchTestOptions.output_dir, fileName,
						//(extractor.getUnit().getPackage() == null? null : extractor.getUnit().getPackage().getName().toString()),
						this.outputPackageName,
						ASTUtils.importDeclarations(extractor.getUnit()), false);
			    files.add(f);
			    //output the next file
			    startIndex += this.testPerFile;
			    
			}
		}
		if(this.outputRandoopTest) {
			GenInputsAbstract.junit_output_dir = SketchTestOptions.output_dir;
			SharedSequencePool.output_randoop_tests();
		}
		return files;
	}
	
	/**
	 * A helper method to check whether all assembled sequences have the
	 * same size
	 * */
	private void checkEqual(List<List<AssembledSequence<?>>> lists) {
		for(int i = 0; i < lists.size() - 1; i++) {//List<AssembledSequence<?>> list : lists) {
			Checker.checkTrue(lists.get(i).size() == lists.get(i+ 1).size(), "Size is not the same, " +
					" the: " + i + " is: " + lists.get(i).size() + ", and the i+ 1 is: " + lists.get(i + 1).size());
		}
	}
	
	private Method getPublicVoidNoArgMethod(String name, Class<?> clazz) {
		Method m = null;
		for(Method method : clazz.getDeclaredMethods()) {
			String methodName = method.getName();
			if(methodName.equals(name) && Modifier.isPublic(method.getModifiers())
					&& !Modifier.isStatic(method.getModifiers())
					&& method.getParameterTypes().length == 0
					&& method.getReturnType().equals(void.class)) {
				m = method;
				break;
			}
		}
		return m;
	}
	
	/**
	 * Check the junit method is a valid one
	 * */
	private void  checkJUnitMethod(Class<?> clazz) {
		Class<?> super_class = clazz.getSuperclass();
		Checker.checkTrue(super_class != null && super_class.getName().equals("junit.framework.TestCase"),
				"The super_class is: " + super_class);
		//check there is no setUp/ tearDown, all method should be test method
//		Method[] methods = clazz.getDeclaredMethods();
//		for(Method method : methods) {
//			String methodName = method.getName();
//			
//			Checker.checkTrue(methodName.startsWith("test"), "All method in clazz : " + clazz.getName() + " should be JUnit test method," +
//					" currently, no setUp, tearDown, and auxiliary methods are supported.");
//			Checker.checkTrue(method.getParameterTypes().length == 0, "The argument length should be 0, for method: " + methodName);
//			Checker.checkTrue(method.getReturnType().equals(void.class), "The return type should be void for method: " + methodName);
//			Checker.checkTrue(Modifier.isPublic(method.getModifiers()), "It should be public for method: " + methodName);
//			Checker.checkTrue(!Modifier.isStatic(method.getModifiers()), "It should not be static for method: " + methodName);
//		}
	}
}